import argparse
import os
import subprocess


def rmdir(path):
    if os.path.isdir(path):
        for file in os.listdir(path):
            rmdir(os.path.join(path, file))
        if os.path.exists(path):
            os.rmdir(path)
    else:
        if os.path.exists(path):
            os.remove(path)


class FileZip(object):
    def __init__(self, config):
        self.config = config
        self.path = config.PATH
        self.zip_bin = config.BIN
        self.father_path = os.path.dirname(self.path)

    def run(self):
        if self.config.MODE == "encrypt":
            fc.encrypt(self.config.PASSWORD)
        elif self.config.MODE == "decrypt":
            fc.decrypt(self.config.PASSWORD)
        else:
            print("mode config error!")

    @staticmethod
    def run_script(cmd):
        p = subprocess.Popen(cmd, close_fds=True, stdout=subprocess.PIPE,
                             stderr=subprocess.STDOUT)
        status = True
        for line in iter(p.stdout.readline, b""):
            try:
                try:
                    msg = str(line, encoding="utf8").strip("\r").strip("\n")
                except:
                    msg = str(line, encoding="gbk").strip("\r").strip("\n")
            except Exception as e:
                msg = str(e)
            print(msg)
            if "ERROR: Wrong password" in msg:
                status = False

        return status

    def encrypt(self, pwd=None):

        zip_name = f"{os.path.basename(self.path)}.zip"
        zip_full_path = os.path.join(self.father_path, zip_name)
        if os.path.exists(zip_full_path):
            print(f"{zip_full_path} exist!")
            return
        cmd = f"{self.zip_bin} a {zip_full_path} {self.path}"
        if pwd:
            cmd += f" -p{pwd}"
        print(cmd)
        en_status = self.run_script(cmd)
        if not en_status:
            print("encrypt failed!")
            if os.path.exists(zip_full_path):
                rmdir(zip_full_path)
                print(f"remove {zip_full_path}")
        if self.config.DEL:
            print(f"delete source file [{self.path}]")
            rmdir(self.path)

    def decrypt(self, pwd=None, dest_path=None):
        if not dest_path:
            dest_path = self.father_path
        base_name = os.path.basename(os.path.splitext(self.path)[0])
        unzip_path = os.path.join(dest_path, base_name)
        if os.path.exists(unzip_path):
            print(f"{dest_path} exist!")
            return
        cmd = f"{self.zip_bin} x {self.path} -o{dest_path}"
        if pwd:
            cmd += f" -p{pwd}"
        de_status = self.run_script(cmd)
        if not de_status:
            print("decrypt failed!")
            if os.path.exists(unzip_path):
                rmdir(unzip_path)
                print(f"remove {unzip_path}")
        if self.config.DEL:
            print(f"delete source file [{self.path}]")
            rmdir(self.path)



def get_arg_parser():
    parser = argparse.ArgumentParser(
        description=(
            'Enter a password to encrypt the compressed folder '
        ),
        add_help=False
    )
    parser.add_argument('PATH', help='file path')
    parser.add_argument(
        '-h',
        '--help',
        action='help',
        help='ceason'
    )

    parser.add_argument(
        '-m',
        action='store',
        dest='MODE',
        choices=['encrypt', 'decrypt'],
        type=str,
        help='mode'
    )

    parser.add_argument(
        '-p',
        action='store',
        dest='PASSWORD',
        default="123456",
        type=str,
        help='password'
    )

    parser.add_argument(
        '-bin',
        action='store',
        dest='BIN',
        default="7z.exe",
        type=str,
        help='7zip bin'
    )

    parser.add_argument(
        '-d',
        action='store',
        dest='DEL',
        choices=[0, 1],
        default=0,
        type=int,
        help='delete source file'
    )
    return parser


if __name__ == '__main__':
    args = get_arg_parser().parse_args()
    print(args)

    if not os.path.exists(args.PATH):
        print("file not found")
    else:
        fc = FileZip(args)
        fc.run()
